import concurrent.futures
import hashlib
import logging
import shutil
from datetime import datetime
from pathlib import Path
from urllib.parse import urlparse

import m3u8
import requests
from tqdm import tqdm

import transcode  # ffmpeg转码

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


# ------- 工具 ------- #
class CommonHelper:
    @staticmethod
    def md5_hash(input_string):
        """
        对给定的字符串进行MD5加密。

        该方法通过hashlib库的md5函数对输入的字符串进行MD5散列计算，返回散列后的16进制字符串。
        主要用于数据完整性校验或密码存储等场景。

        参数:
            input_string (str): 需要进行MD5散列的字符串。

        返回:
            str: 输入字符串的MD5散列值，为一个32位的16进制字符串。
        """
        # 初始化MD5加密对象
        md5 = hashlib.md5()
        # 将输入字符串转换为UTF-8编码并更新MD5加密对象
        md5.update(input_string.encode('utf-8'))
        # 返回加密后的16进制字符串
        return md5.hexdigest()

    @staticmethod
    def print_time(message):
        """
        打印当前时间戳和指定信息。

        通过静态方法打印出当前时间戳和指定的信息，方便在日志中记录事件发生的时间。
        这对于调试和追踪程序运行过程中的问题非常有用。

        参数:
        message: str - 需要打印的信息。
        """
        # 获取当前时间并格式化为字符串
        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        # 使用logger记录信息和当前时间
        logger.debug(f"{message} {current_time}")

    @staticmethod
    def get_url_last_path(url):
        """
        获取URL的最后路径部分。

        该方法用于解析一个URL字符串，并返回其路径部分的最后一个元素。
        这对于处理链接并获取其指向的具体资源名非常有用。

        参数:
            url (str): 需要解析的URL字符串。

        返回:
            str: URL路径部分的最后一个元素。
        """
        # 解析URL并获取路径部分
        path = urlparse(url).path
        # 将路径字符串按斜杠分割成列表，并返回列表的最后一个元素
        return path.split('/')[-1]


class M3u8Downloader(CommonHelper):
    """
    M3u8下载器类，继承自CommonHelper类。
    用于下载m3u8格式的视频文件，并可并发下载文件。
    """

    def __init__(self, callback, max_repeat_times=3, download_path='./downloads'):
        """
        初始化M3u8Downloader对象。

        :param callback: 下载完成后的回调函数。
        :param max_repeat_times: 文件下载失败时的最大重试次数，默认为3。
        :param download_path: 文件下载的保存路径，默认为"./downloads"。
        """
        # 初始化下载统计信息
        self.downloaded_count = 0
        self.total_count = 0
        self.failed_count = 0
        # 初始化重试次数字典
        self.repeat_times_dict = {}
        # 设置最大重试次数
        self.max_repeat_times = max_repeat_times
        # 设置下载路径
        self.download_path = Path(download_path)
        self.default_download_path = self.download_path
        # 初始化文件名
        self.file_name = ''
        # 设置回调函数
        self.callback = callback
        # 设置代理
        self.proxies = None
        # 初始化requests会话
        self.session = requests.Session()
        # 禁用urllib3的警告
        requests.packages.urllib3.disable_warnings()
        # 设置http和https的连接池大小
        self.session.mount('http://', requests.adapters.HTTPAdapter(pool_connections=50, pool_maxsize=50))
        self.session.mount('https://', requests.adapters.HTTPAdapter(pool_connections=50, pool_maxsize=50))

    @classmethod
    def parse_m3u8(cls, m3u8_url):
        """
        解析m3u8链接，返回最深层的playlist。

        :param m3u8_url: m3u8文件的URL。
        :return: 解析后的playlist对象。
        """
        playlist = m3u8.load(m3u8_url)
        if playlist.is_variant:
            if playlist.playlists:
                sub_playlist = playlist.playlists[0]
                return cls.parse_m3u8(sub_playlist.absolute_uri)
        else:
            return playlist

    @staticmethod
    def validate_parameter(param, is_url):
        """
        验证参数是否有效。

        :param param: 待验证的参数。
        :param is_url: 参数是否为URL。
        :return: 参数是否有效的布尔值。
        """
        if isinstance(param, str):
            if is_url:
                return param.startswith(('http://', 'https://')) and len(param) > 15
            else:
                return len(param) > 3
        return False

    def get_file_urls_and_paths(self, playlist):
        """
        获取所有文件的URL和保存路径。

        :param playlist: m3u8的playlist对象。
        :return: 文件URL列表和文件保存路径列表。
        """
        all_urls = []
        all_file_paths = []
        temp_path = self.temp_path(self.download_path, self.file_name)
        for seg in filter(None, playlist.segments):
            all_urls.append(seg.absolute_uri)
            fn = self.md5_hash(self.get_url_last_path(seg.absolute_uri))
            all_file_paths.append(temp_path / fn)

        for key in filter(None, playlist.keys):
            all_urls.append(key.absolute_uri)
            fn = self.md5_hash(self.get_url_last_path(key.absolute_uri))
            all_file_paths.append(temp_path / fn)

        return all_urls, all_file_paths

    @classmethod
    def create_native_m3u8_file(cls, playlist, m3u8_path):
        """
        创建一个去除了绝对路径的m3u8文件。

        :param playlist: m3u8的playlist对象。
        :param m3u8_path: 生成的m3u8文件的保存路径。
        """
        for seg in filter(None, playlist.segments):
            seg.uri = cls.md5_hash(cls.get_url_last_path(seg.absolute_uri))

        for key in filter(None, playlist.keys):
            key.uri = cls.md5_hash(cls.get_url_last_path(key.absolute_uri))

        playlist.dump(str(m3u8_path))

    def download_file(self, url, file_path):
        """
        下载单个文件。

        :param url: 文件的URL。
        :param file_path: 文件的保存路径。
        """
        file_name = file_path.name
        if file_path.exists():
            self.downloaded_count += 1
            logger.debug(f"文件 {file_name} 已存在({self.downloaded_count}/{self.total_count})")
            return
        response = self.session.get(url, verify=False)
        # response = requests.get(url, verify=False, proxies=self.proxies)
        if response.status_code == 200:
            with file_path.open('wb') as file:
                file.write(response.content)

            self.downloaded_count += 1
            logger.debug(f"文件 {file_name} 下载完成({self.downloaded_count}/{self.total_count})")

        else:
            retry_times = self.repeat_times_dict.get(url, 0)
            logger.error(f"无法下载文件 {file_name}: {response.status_code}")
            if retry_times >= self.max_repeat_times:
                logger.debug(f'文件 {file_name} 已达到最大下载次数 {self.max_repeat_times} 次')
                logger.debug(f'文件url:{url}')
                self.failed_count += 1
            else:
                self.repeat_times_dict[url] = retry_times + 1
                logger.debug(f'开始第{retry_times + 1}次重试')
                self.download_file(url, file_path)

    def download_files_concurrently(self, url_list, filename_list):
        """
        并发下载文件。

        :param url_list: 文件URL列表。
        :param filename_list: 文件保存路径列表。
        """
        with concurrent.futures.ThreadPoolExecutor(max_workers=50) as executor:
            tasks = [executor.submit(self.download_file, url, filename) for url, filename in
                     zip(url_list, filename_list)]

            # 使用tqdm的上下文管理器来更新进度条
            for future in tqdm(concurrent.futures.as_completed(tasks), total=len(url_list), desc="Downloading"):
                result = future.result()
                if result is not None:
                    logger.info("Task completed:", result)

            logger.info("All tasks are done.")
            self.download_complete()

    @staticmethod
    def temp_path(download_path, file_name):
        """
        创建一个临时文件夹。

        :param download_path: 下载路径。
        :param file_name: 文件名。
        :return: 临时文件夹的路径。
        """
        temp_path = download_path / file_name
        temp_path.mkdir(parents=True, exist_ok=True)
        return temp_path

    @classmethod
    def m3u8_path(cls, download_path, file_name):
        """
        获取m3u8文件的保存路径。

        :param download_path: 下载路径。
        :param file_name: 文件名。
        :return: m3u8文件的保存路径。
        """
        return cls.temp_path(download_path, file_name) / 'index.m3u8'

    @staticmethod
    def mp4_path(download_path, file_name):
        """
        获取最终生成的mp4文件的保存路径。

        :param download_path: 下载路径。
        :param file_name: 文件名。
        :return: mp4文件的保存路径。
        """
        return download_path / f"{file_name.rstrip('.mp4')}.mp4"

    def download_m3u8(self, m3u8_url, file_name, download_path=''):
        """
        下载m3u8文件，并将其转码为mp4文件。

        :param m3u8_url: m3u8文件的URL。
        :param file_name: 下载的文件名。
        :param download_path: 下载保存的路径。
        """
        self.download_path = Path(download_path) if download_path else self.default_download_path
        if not self.validate_parameter(m3u8_url, True):
            logger.error('m3u8的url不正确，请重新输入')
            return
        if not self.validate_parameter(file_name, False):
            logger.error('请输入正确的文件名')
            return

        self.file_name = file_name
        mp4_path = self.mp4_path(self.download_path, self.file_name)
        if mp4_path.exists():
            logger.info(f'[{file_name}] 文件已存在！！！')
            return

        logger.info('开始获取m3u8的内容')
        logger.info(f'm3u8URL: {m3u8_url}')
        playlist = self.parse_m3u8(m3u8_url)
        file_urls, file_paths = self.get_file_urls_and_paths(playlist)
        logger.info('解析m3u8完成，得到所有的文件url')

        m3u8_path = self.m3u8_path(self.download_path, self.file_name)
        self.create_native_m3u8_file(playlist, m3u8_path)

        self.downloaded_count = 0
        self.failed_count = 0
        self.total_count = len(file_urls)

        if file_urls and file_paths:
            logger.info('开始下载所有文件(ts和key)')
            self.download_files_concurrently(file_urls, file_paths)
        else:
            self.print_time(f'文件 {file_name} 下载失败!')

    def download_complete(self):
        """
        检查下载任务是否完成，并进行转码操作。

        如果下载失败的文件数量大于0，则打印下载失败的消息；
        否则，打印所有文件下载完成的消息，并开始转码过程。

        转码完成后，调用回调函数并删除临时文件。
        如果转码失败，则调用回调函数传入None。
        """
        # 检查是否有文件下载失败
        if self.failed_count > 0:
            self.print_time(f'下载失败！有{self.failed_count}个文件未能下载，请重试！！')
        else:
            self.print_time(f'所有文件已下载完成！！！下面开始转码！！！')

        # 构建转码前后的文件路径
        m3u8_path = self.m3u8_path(self.download_path, self.file_name)
        mp4_path = self.mp4_path(self.download_path, self.file_name)
        temp_path = self.temp_path(self.download_path, self.file_name)

        # 尝试将m3u8文件转码为mp4
        if transcode.m3u8_to_mp4(str(m3u8_path), str(mp4_path)):
            self.print_time(f'文件 [{self.file_name}] 转码完成！')
            # 转码成功后，调用回调函数并删除临时文件
            self.callback(str(mp4_path))
            shutil.rmtree(str(temp_path))
            self.print_time("已清空临时文件！！")
        else:
            self.print_time(f'文件 [{self.file_name}] 转码失败！')
            # 转码失败，调用回调函数传入None
            self.callback(None)
